import os
import numpy as np
from sklearn.datasets import load_files
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.model_selection import train_test_split
from sklearn.naive_bayes import MultinomialNB
from sklearn.metrics import accuracy_score

class TextClassifier:
    def __init__(self, alpha=1.0):
        self.alpha = alpha
        self.model = None
        self.vectorizer = None

    def fit(self, X, y):
        # 使用TF-IDF向量化文本数据
        self.vectorizer = TfidfVectorizer(stop_words='english', max_features=1000)
        X_transformed = self.vectorizer.fit_transform(X)
        
        # 初始化并训练多项式朴素贝叶斯模型
        self.model = MultinomialNB(alpha=self.alpha)
        self.model.fit(X_transformed, y)

    def predict(self, X):
        X_transformed = self.vectorizer.transform(X)
        return self.model.predict(X_transformed)

def load_local_data(local_path, categories=None):
    if not os.path.exists(local_path):
        raise FileNotFoundError(f"The specified path does not exist: {local_path}")

    print("Loading data from local path...")
    newsgroups = load_files(local_path, categories=categories, encoding='latin-1')
    X = newsgroups.data
    y = newsgroups.target
    return train_test_split(X, y, test_size=0.2, random_state=42)

if __name__ == "__main__":
    try:
        # 指定本地数据集路径
        local_data_path = "./20_newsgroups"
        categories = ['sci.space', 'rec.autos']  # 只选择两个类别进行测试

        # 加载数据
        X_train, X_test, y_train, y_test = load_local_data(local_data_path, categories=categories)

        # 初始化并训练模型
        classifier = TextClassifier(alpha=1.0)
        classifier.fit(X_train, y_train)

        # 预测并计算准确率
        y_pred = classifier.predict(X_test)
        accuracy = accuracy_score(y_test, y_pred)
        print(f"测试集准确率: {accuracy:.4f}")
    except Exception as e:
        print(f"An error occurred: {e}")